#include "stdafx.h"
#include "ETHFram.h"
#include "TCPSession.h"

#pragma comment(lib, "ws2_32.lib")
using namespace std;

vector<IN_ADDR> CETHFram::gLocalAdds;

CETHFram::CETHFram(const unsigned char*data, unsigned int plen) :mDatalen(plen), mPreProcessFlag(false), mEthHdr(NULL),
mIpHdr(0), mTcpHdr(0), mUdpHdr(0), mIPProtocol(0), mHdrlen(0), mPackettotallen(0), mTCPPayloadlen(0)
{
	mData = new unsigned char[plen];
	if (mData)
		memcpy_s(mData, plen, data, plen);
}


CETHFram::~CETHFram()
{
	if (mData) {
		delete[] mData;
		mData = NULL;
	}
}

bool CETHFram::Initialize()
{
	WSADATA wsaData;
	int iResult;
	iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (iResult != 0) {
		printf("WSAStartup failed: %d\n", iResult);
		return 1;
	}

	char ac[MAX_PATH] = { 0 };
	if (gethostname(ac, sizeof(ac)) == SOCKET_ERROR) {
		return false;
	}

	struct hostent *phe = gethostbyname(ac);
	if (phe == 0) {
		return false;
	}

	for (int i = 0; phe->h_addr_list[i] != 0; ++i) {
		struct in_addr addr;
		memcpy(&addr, phe->h_addr_list[i], sizeof(struct in_addr));
		gLocalAdds.push_back(addr);
	}
	return true;
}

bool CETHFram::VerifyIPAddr()
{
	if (!mIpHdr)
		return false;
	for (vector<IN_ADDR>::iterator it = gLocalAdds.begin(); it != gLocalAdds.end(); it++) {
		IN_ADDR local = *it;
		if (memcmp(&mIpHdr->ip_dst, &local, sizeof(IN_ADDR))==0 || memcmp(&mIpHdr->ip_src, &local, sizeof(IN_ADDR)) == 0)
			return true;
	}
	return false;
}

INT32 CETHFram::Process()
{
	CTcpSession* session = nullptr;
	if (mTcpHdr->th_flags & TH_FIN || mTcpHdr->th_flags & TH_RST) {
		string sip = inet_ntoa(mIpHdr->ip_dst);
		string dip = inet_ntoa(mIpHdr->ip_src);
		//printf("Close: ip.addr == %s && ip.addr == %s\n", sip.c_str(), dip.c_str());
		CTCPSessionMgr::CloseSession(mIpHdr, mTcpHdr);
		return 0;
	}
	if (mTcpHdr->th_flags & TH_SYN && !(mTcpHdr->th_flags & TH_ACK)) {
		string sip = inet_ntoa(mIpHdr->ip_dst);
		string dip = inet_ntoa(mIpHdr->ip_src);
		//printf("Open: ip.addr == %s && ip.addr == %s && tcp.port == %u && tcp.port == %u\n", sip.c_str(), dip.c_str(), ntohs(mTcpHdr->th_sport), ntohs(mTcpHdr->th_dport));
		session = CTCPSessionMgr::AddSession(mIpHdr, mTcpHdr);
		return 0;
	}
	else {
		session = CTCPSessionMgr::GetSession(mIpHdr, mTcpHdr);
	}

	if (!session) {
		return -1;
	}

	if (mPackettotallen + sizeof(ethhdr_t) - mHdrlen <= 0) {
		return -1;
	}
	mTCPPayloadlen = mDatalen - mHdrlen;

	return session->HandlePacket(this, mIpHdr, mTcpHdr, mData + mHdrlen, mDatalen - mHdrlen);
}

UINT32 CETHFram::GetHashVal(UINT32 width)
{
	if (PreProcess() == false) {
		return -1;
	}
	
	return (mTcpHdr->th_dport&mTcpHdr->th_sport) % width ;
}

bool CETHFram::PreProcess()
{
	if (mPreProcessFlag) {
		return true;
	}
	mEthHdr = reinterpret_cast<const ethhdr_t*>(mData);

	if (htons(mEthHdr->h_proto) != ETH_P_IP) {
		return false;
	}

	mHdrlen += sizeof(ethhdr_t);
	mIpHdr = reinterpret_cast<const iphdr_t*>(mData + sizeof(ethhdr_t));
	if (mIpHdr->ip_v != IPVERSION)
		return false;

#if !DEBUGFILE
	if (VerifyIPAddr() == false) {
		return false;
	}
#endif

	mPackettotallen = ntohs(mIpHdr->ip_len);
	mHdrlen += mIpHdr->ip_hl << 2;
	if (mIpHdr->ip_p == IPPROTO_TCP) {
		mTcpHdr = reinterpret_cast<const tcphdr_t*>(mData + mHdrlen);
		mHdrlen += mTcpHdr->th_off << 2;
	}
	else {
		return false;
	}

	mPreProcessFlag = true;
	return true;
}
